#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   eth.py
@Time    :   2024-08-23 09:17
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   批量管理lxc容器网络接口
"""

import logging
import sys
from argparse import ArgumentParser
from subprocess import getstatusoutput
from sys import exit

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


class CtEthDelete:
	def __init__(self):
		self.net_list = []
		self.net_dict = {}
		self.net_type = ""
		self.net_delete_start = 0
		self.net_delete_end = 0
		self.host_id_dict = {}
		self.host_eth_count = 0
		self.id = None
		self.net_count = 0
		self.interface_dict = {}
		self.interface_name = 'vmbr0'

	def get_host_id_list(self):
		"""
		获取容器ID列表
		:return:
		"""
		status, result = getstatusoutput("pct list  | awk '{print $1,$3}' | grep -v VMID")
		if status == 0:
			for i in result.split("\n"):
				self.host_id_dict[i.split(" ")[0]] = i.split(" ")[1]
		else:
			print("获取容器ID列表失败")
			exit(1)

	def select_host(self):
		"""
		选择需要操作的容器ID
		:return:
		"""
		while True:
			print("请选择需要操作的容器ID：")
			for i in self.host_id_dict:
				print(f"{i} {self.host_id_dict[i]}")
			print("0. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "0":
				exit(0)
			elif choice in self.host_id_dict:
				self.id = choice
				break
			else:
				print("输入错误，请重新输入")

	def get_host_eth_nmber(self):
		"""
		获取容器网卡数量
		:return:
		"""
		cmd = f'pct config {self.id} | grep ^net | wc -l'
		status, net_id = getstatusoutput(cmd)
		if status == 0:
			print(cmd)
			logger.info(f"当前容器{self.id}已存在{net_id}个网卡")
			self.host_eth_count = int(net_id)

	def delete_eth(self):
		"""
		开始添加网卡
		:return:
		"""
		cmd_list = []
		id_ = self.host_eth_count + 1
		logger.info("欲执行的命令如下：")
		for i in range(self.net_delete_start, self.net_delete_end + 1):
			cmd = f"pct set {self.id} --delete net{i}"
			print(cmd)
			id_ += 1
			cmd_list.append(cmd)
		# 询问是否开始添加
		while True:
			print("是否开始删除网卡")
			print("1. 开始删除")
			print("2. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "1":
				n_ = 1
				for i in cmd_list:
					status, result = getstatusoutput(i)
					if status == 0:
						logger.info(f"删除第 [{n_} ] 张网卡成功")
					else:
						logger.error(f"删除第 [{n_} ] 张网卡失败")
						print(result)
					n_ += 1
				break
			elif choice == "2":
				exit(0)
				break
			else:
				print("输入错误，请重新输入")

	def set_start(self):
		"""
		设置网卡数量
		:return:
		"""
		while True:
			logger.info("当前容器已存在" + str(self.host_eth_count) + "个网卡")
			try:
				self.net_delete_start = int(input("请设置需要删除的网卡起始ID：\n"))
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break

	def set_end(self):
		"""
		设置网卡数量
		:return:
		"""
		while True:
			logger.info("当前容器已存在" + str(self.host_eth_count) + "个网卡")
			logger.info(f"当前已设置起始网卡ID为：" + str(self.net_delete_start))
			try:
				self.net_delete_end = int(input("请设置需要删除的网卡结束ID：\n"))
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break

	def start(self):
		"""
		开始配置过程，按顺序执行一系列配置步骤。

		此方法无参数，通过调用不同方法实现容器选择、网络类型选择、数量设置、接口列表获取与设置，
		最终完成添加操作。这些步骤是配置过程中的关键环节，确保按照预定顺序执行。
		"""
		self.get_host_id_list()  # 获取容器ID列表
		self.select_host()  # 选择容器
		self.get_host_eth_nmber()
		self.set_start()  # 设置起始
		self.set_end()  # 设置结束
		self.delete_eth()  # 开始删除


class CtEthAdd:
	def __init__(self, arg: dict):
		"""

		:param arg:
		"""

		self.arg_end = arg['end']
		self.arg_start = arg['start']
		# self.arg_dns2 = arg['dns2']
		# self.arg_dns1 = arg['dns1']
		self.arg_sub = arg['sub']
		self.arg_gw = arg['gateway']
		self.arg_netmask = arg['netmask']
		self.arg_br = arg['br']
		self.id = arg['id']
		self.arg_ips = arg['ips']
		self.net_list = []
		self.net_dict = {}
		self.net_type = ""
		self.host_id_dict = {}
		self.host_eth_count = 0
		# self.id = None
		self.net_count = int(int(self.arg_end) - int(self.arg_start) + 1)
		self.interface_dict = {}
		self.interface_name = 'vmbr0'

	def select_net_type(self):
		"""
		选择pve虚拟机网卡类型
		:return:
		"""
		while True:
			print("请选择网卡类型：")
			print("1. [virtio] 默认类型,半虚拟化,拥有最接近物理的性能,但是部分系统无驱动,无法显示接口速率(不影响使用)")
			print("2. [intel e1000] 兼容模式,相比性能会稍微慢,但是兼容性更好,可以显示接口速率(1000M)")
			print("3. [Realtek RTL8139] 兼容模式,兼容性更好,可以显示接口速率(100M)")
			print(
				"4. [intel e1000e] 兼容模式,相比性能会稍微慢,但是兼容性更好,可以显示接口速率(10000M),PVE版本需大于8.1")
			print("0. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "0":
				exit(0)
			elif choice == "1":
				self.net_type = "virtio"
				break
			elif choice == "2":
				self.net_type = "e1000"
				break
			elif choice == "3":
				self.net_type = "rtl8139"
				break
			elif choice == "4":
				self.net_type = "e1000e"
				break
			else:
				print("输入错误，请重新输入")

	def get_host_id_list(self):
		"""
		获取容器ID列表
		:return:
		"""
		status, result = getstatusoutput("pct list  | awk '{print $1,$3}' | grep -v VMID")
		if status == 0:
			for i in result.split("\n"):
				self.host_id_dict[i.split(" ")[0]] = i.split(" ")[1]
		else:
			print("获取容器ID列表失败")
			exit(1)

	def select_host(self):
		"""
		选择需要操作的容器ID
		:return:
		"""
		if self.id is None:
			while True:
				print("请选择需要操作的容器ID：")
				for i in self.host_id_dict:
					print(f"{i} {self.host_id_dict[i]}")
				print("0. 退出")
				choice_ = input("请输入您的选择：\n")
				if choice_ == "0":
					exit(0)
				elif choice_ in self.host_id_dict:
					self.id = choice_
					break
				else:
					print("输入错误，请重新输入")
		else:
			if self.id in self.host_id_dict:
				logger.info(f"您选择的容器ID为：{self.id}")
			else:
				logger.error("您选择的容器ID不存在")
				exit(1)

	def get_host_eth_nmber(self):
		"""
		获取容器网卡数量
		:return:
		"""
		cmd = f'pct config {self.id} | grep net | grep name | wc -l'
		status, net_id = getstatusoutput(cmd)
		if status == 0:
			print(cmd)
			logger.info(f"当前容器{self.id}已存在{net_id}个网卡")
			self.host_eth_count = int(net_id)

	def add_list(self):
		"""

		:return:
		"""
		cmd_list = []
		id_ = self.host_eth_count + 1
		logger.info("欲执行的命令如下：")
		for i in str(self.arg_ips).split(' '):
			ip_ = f"{self.arg_sub}.{i}"
			cmd = f"pct set {self.id} --net{id_} name=eth{id_},bridge={self.interface_name},ip={ip_}/{self.arg_netmask},gw={self.arg_gw}"
			print(cmd)
			id_ += 1
			cmd_list.append(cmd)
		# 询问是否开始添加
		while True:
			print("是否开始添加网卡？")
			print("1. 开始添加")
			print("2. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "1":
				n_ = 1
				for i in cmd_list:
					status, result = getstatusoutput(i)
					if status == 0:
						logger.info(f"添加第 [{n_} ] 张网卡成功")
					else:
						logger.error(f"添加第 [{n_} ] 张网卡失败")
					n_ += 1
			elif choice == "2":
				exit(0)
				break
			else:
				print("输入错误，请重新输入")

	def add(self):
		"""
		开始添加网卡
		:return:
		"""
		cmd_list = []
		id_ = self.host_eth_count + 1
		logger.info("欲执行的命令如下：")
		for i in range(self.arg_start, self.arg_end + 1):
			ip_ = f"{self.arg_sub}.{i}"
			cmd = f"qm set {self.id} --net{id_} {self.net_type},bridge={self.interface_name}"
			cmd = f"pct set {self.id} --net{id_} name=eth{id_},bridge={self.interface_name},ip={ip_}/{self.arg_netmask},gw={self.arg_gw}"
			print(cmd)
			id_ += 1
			cmd_list.append(cmd)
		# 询问是否开始添加
		while True:
			print("是否开始添加网卡？")
			print("1. 开始添加")
			print("2. 退出")
			choice = input("请输入您的选择：\n")
			if choice == "1":
				n_ = 1
				for i in cmd_list:
					status, result = getstatusoutput(i)
					if status == 0:
						logger.info(f"添加第 [{n_} ] 张网卡成功")
					else:
						logger.error(f"添加第 [{n_} ] 张网卡失败")
					n_ += 1
			elif choice == "2":
				exit(0)
				break
			else:
				print("输入错误，请重新输入")

	def get_interface_list(self):
		"""
		获取网卡接口列表
		:return:
		"""
		id_ = 1
		status, result = getstatusoutput("ip link | awk '{print $2}' | sed 's/://' | grep vmbr")
		if status == 0:
			for i in result.split("\n"):
				self.interface_dict[id_] = i
				logger.info(f"{id_} {i}")
				id_ += 1
		else:
			print("获取网卡接口列表失败")
			exit(1)

	def set_interface(self):
		"""
		设置网卡接口
		:return:
		"""
		if self.arg_br is not None:
			self.interface_name = self.arg_br
			if self.interface_name not in self.interface_dict.values():
				logger.warning("输入的网卡接口不存在")
				exit(1)
			else:
				return
		while True:
			print("请设置网卡接口：")
			for i in self.interface_dict:
				print(f"{i} {self.interface_dict[i]}")
			try:
				self.interface_name = self.interface_dict[int(input("请输入您的选择："))]
				break
			except ValueError:
				logger.warning("输入错误，请重新输入")

	def set_number(self):
		"""
		设置网卡数量
		:return:
		"""
		if self.arg_ips is not None:
			self.net_count = len(str(self.arg_ips).split(' ')) - self.host_eth_count
			return
		res = 'n'
		while True:
			logger.info("当前容器已存在" + str(self.host_eth_count) + "个网卡")
			logger.info(f"当前设置新增网卡数量为" + str(self.net_count) + "张")
			try:
				res = str(input("是否继续[y/n]：\n")).strip()
			except ValueError:
				print("输入错误，请重新输入")
			else:
				break
			if res == 'y':
				logger.info("继续设置网卡数量")
			else:
				logger.info("退出设置网卡")
				exit(0)

	def start(self):
		"""
		开始配置过程，按顺序执行一系列配置步骤。

		此方法无参数，通过调用不同方法实现容器选择、网络类型选择、数量设置、接口列表获取与设置，
		最终完成添加操作。这些步骤是配置过程中的关键环节，确保按照预定顺序执行。
		"""
		self.get_host_id_list()  # 获取容器ID列表
		self.select_host()  # 选择容器
		self.get_host_eth_nmber()  # 获取网卡数量
		self.select_net_type()  # 选择网络类型
		self.set_number()  # 设置数量
		self.get_interface_list()  # 获取接口列表
		self.set_interface()  # 设置接口
		if len(str(self.arg_ips)) > 1:
			self.add_list()
		else:
			self.add()  # 添加配置


def get_default_gw():
	"""
	获取默认网关
	:return:
	"""
	gw = getstatusoutput('ip route | grep default')
	if gw[0] == 0:
		gw = gw[1].split(' ')[2]
		return gw
	else:
		return None


if __name__ == '__main__':
	ver = '1.0'
	sub_ = '192.168.1'
	gw_ = '192.168.1.1'
	gw = get_default_gw()
	if gw is None:
		logger.warning('获取默认网关失败')
		logger.warning("Failed to obtain default gateway")
	else:
		sub_ = f"{gw.split('.')[0]}.{gw.split('.')[1]}.{gw.split('.')[2]}"
		gw_ = gw
	parser = ArgumentParser(description=f'批量管理CT容器网卡并设置静态IP: {ver} 作者: 坐公交也用券')
	parser.add_argument('-s', '--sub', help=f'网段,默认：{sub_}', required=False, default=sub_)
	parser.add_argument('-g', '--gateway', help=f'网关,默认: {gw_}', required=False, default=gw_)
	parser.add_argument('-n', '--netmask', help='子网掩码,默认: 24', required=False, default=24, type=int)
	parser.add_argument('-start', help='网卡起始序号,默认: 1', required=False, default=1, type=int)
	parser.add_argument('-end', help='网卡结束序号,默认: 1', required=False, default=1, type=int)
	parser.add_argument('-id', help='容器/虚拟机ID,如果未传入此参数则从键盘输入', required=False)
	parser.add_argument('-br', help='网卡接口名称,例如：vmbr0,如果未传入此参数则从键盘输入', required=False)
	parser.add_argument('-ips', help='设置批量网卡IP地址(传入此参数后,start和end参数将失效),每个IP之间使用空格间隔,参数使用双引号,例如："2 3 4"',
	                    required=False)
	args = parser.parse_args()
	args_dict = vars(args)
	print(args_dict)
	logger.info(f"版本: {ver}")
	while True:
		choice = input("请选择要执行的操作：\n1. 添加网卡\n2. 删除网卡\n0. 退出\n")
		if choice == "1":
			logger.info("添加网卡")
			n = CtEthAdd(arg=args_dict)
			n.start()
		elif choice == "2":
			logger.info("删除网卡")
			n = CtEthDelete()
			n.start()
		elif choice == "0":
			exit(0)
		else:
			logger.warning("输入错误，请重新输入")
